#！ /usr/bin/env python
# -*- coding: utf-8 -*-
# Date:

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import tushare as ts
import datetime
import dateutil

ts.set_token('你申请的token')
pro = ts.pro_api()

# df = pro.daily_basic(ts_code='', trade_date='20180726', fields='ts_code,trade_date,turnover_rate,volume_ratio,pe,pb')
# 该接口需要2000积分才能使用。
# print(df)

CASH = 100000
START_DATE = '20140101'
END_DATE = '20150101'

trade_cal = pro.trade_cal()
trade_cal = trade_cal.sort_values(by=['cal_date'], ascending=True)
class Context():
    def __init__(self, cash, start_date, end_date):
        self.cash = cash
        self.start_date = start_date
        self.end_date = end_date
        self.positions = {}
        self.benchmark = None
        self.date_range = trade_cal[(trade_cal['is_open'] == 1) & \
                                    (trade_cal['cal_date'] >= start_date) & \
                                    (trade_cal['cal_date'] <= end_date)]['cal_date'].values
        self.dt = None

context = Context(CASH, START_DATE, END_DATE)

class G:
    pass

g = G()

def set_benchmark(security):        # 只支持一只股票作为基准
    context.benchmark = security

def attribute_history(security, count, fields=('trade_date','open','close','high','low','vol')):
    end_date = (context.dt - datetime.timedelta(days=1)).strftime("%Y%m%d")
    start_date = trade_cal[(trade_cal['is_open'] == 1) & (trade_cal['cal_date'] <= end_date)][:count].iloc[count-1, :]['cal_date']
    return attribute_daterange_history(security, start_date, end_date, fields)

def attribute_daterange_history(security, start_date, end_date, fields=('trade_date','open','close','high','low','vol')):
    try:
        f = open(security+'.csv', 'r')
        df = pd.read_csv(f , index_col='trade_date', parse_dates=['trade_date']).loc[start_date:end_date, :].sort_values(by=['cal_date'], ascending=True)
        df = df.set_index('cal_date', drop=True)
    except FileNotFoundError:
        df = pro.daily(ts_code=security, start_date=start_date, end_date=end_date).sort_values(by=['trade_date'], ascending=True)
        df = df.set_index('trade_date', drop=False)
    return df[list(fields)]

def get_today_data(security):
    today = context.dt.strftime('%Y%m%d')
    try:
        f = open(security+'.csv', 'r')
        data = pd.read_csv(f, index_col='trade_date', parse_dates=['trade_date']).loc[today,:]
    except FileNotFoundError:
        data = pro.daily(ts_code=security, start_date=today, end_date=today).iloc[0, :]
    except KeyError:    # 停牌
        data = pd.Series()
    return data

# print(get_today_data('601318.SH'))

def _order(today_data, security, amount):
    p = today_data['open']

    if len(today_data) == 0:
        print("今日停牌")
        return

    if context.cash - amount * p < 0:
        amount = int(context.cash / p)
        print("现金不足，已调整为%d" % (amount))

    if amount % 100 != 0:
        if amount != -context.positions.get(security, 0):
            amount = int(amount / 100) * 100
            print("不是100的倍数，已调整为%d" % amount)

    if context.positions.get(security, 0) < -amount:
        amount = -context.positions.get(security, 0)
        print("卖出股票不能超过持仓的数量，已调整为%d" % amount)

    context.positions[security] = context.positions.get(security, 0) + amount

    context.cash -= amount * p

    if context.positions[security] == 0:
        del context.positions[security]


def order(security, amount):
    today_data = get_today_data(security)
    _order(today_data, security, amount)

def order_target(security, amount):
    if amount < 0:
        print("数量不能为负，已调整成0")
        amount = 0

    today_data = get_today_data(security)
    hold_amount = context.positions.get(security, 0)    # ToDo: T + 1 closeable total
    delta_amount = amount - hold_amount
    _order(today_data, security, delta_amount)

def order_value(security, value):
    today_data = get_today_data(security)
    amount = int(value / today_data['open'])
    _order(today_data, security, amount)

def order_target_value(security, value):
    today_data = get_today_data(security)
    if value < 0:
        print("价值不能为负，已调整为0")
        value = 0

    hold_value = context.positions.get(security, 0) * today_data['open']
    delta_value = value - hold_value
    order_value(security, delta_value)


def run():
    plt_df = pd.DataFrame(index=pd.to_datetime(context.date_range).strftime("%Y%m%d"), columns=['value'])
    init_value = context.cash
    initialize(context)
    last_price = {}
    for dt in context.date_range:
        context.dt = dateutil.parser.parse(dt)
        handle_data(context)
        value = context.cash
        for stock in context.positions:
            # 考虑停牌的情况
            today_data = get_today_data(stock)
            if len(today_data) == 0:
                p = last_price[stock]
            else:
                p = get_today_data(stock)['open']
                last_price[stock] = p
            value += p * context.positions[stock]
        plt_df.loc[dt, 'value'] = value
    plt_df['ratio'] = (plt_df['value'] - init_value) / init_value

    bm_df = attribute_daterange_history(context.benchmark, context.start_date, context.end_date)
    bm_init = bm_df['open'][0]
    plt_df['benchmark_ratio'] = (bm_df['open'] - bm_init) / bm_init

    plt_df[['ratio','benchmark_ratio']].plot()
    plt.show()

def initialize(context):
    set_benchmark('601318.SH')
    g.p1 = 5
    g.p2 = 60
    g.security = '601318.SH'

def handle_data(context):
    hist = attribute_history(g.security, g.p2)
    ma5 = hist['close'][-g.p1:].mean()
    ma60 = hist['close'].mean()

    if ma5 > ma60 and g.security not in context.positions:
        order_value(g.security, context.cash)
    elif ma5 < ma60 and g.security in context.positions:
        order_target(g.security, 0)

run()